import torch
import torch.nn as nn
import torch.nn.functional as F
import math

import os
import sys
project_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
sys.path.append(project_dir)

from transcoder.models.bsqvit import VITBSQModel
from transcoder.models.communicator.csi import CSISemanticModel
from transcoder.models.communicator.beamform import BeamformModel
from transcoder.models.communicator.transmit import TransMitModel
from transcoder.models.communicator.demodulate import DemodulateModel
from transcoder.losses.channel_aware import ChannelAwareLoss
from transcoder.losses.semantic_align import SemanticAlignmentLoss

class SCNetModel(nn.Module):
    def __init__(self, bsqconfig, semconfig):
        super().__init__()
        vitconfig = bsqconfig['vitconfig']
        self.K  = semconfig['K']
        self.Nc = semconfig['Nc']
        self.Q  = semconfig['Q']
        self.Nt = semconfig['Nt']
        
        # ========== 图像量化模块 (VITBSQModel) ========== #
        self.image_quantizer = VITBSQModel(**bsqconfig)
        
        # ========== 联合信源信道编码 ========== #
        # MLP适配器 (将BSQ特征映射到子载波维度)
        self.Hadaptor = nn.Sequential(
            nn.Linear(bsqconfig['embed_dim'], 512),
            nn.GELU(),
            nn.Linear(512, self.Nc)
        )
        # QPSK调制，调整载波信号的相位来传输数字信息
        self.patches = (vitconfig['image_size'] // vitconfig['patch_size'])**2
        self.qpsk = nn.Sequential(
            nn.Linear(self.patches, 512),
            nn.GELU(),
            nn.Linear(512, 2*self.Q)  # 输出实部+虚部各Nc
        )
        
        # ========== 通信模块 ========== #
        # CSI特征提取
        self.csi_extractor = CSISemanticModel(**semconfig)
        
        # 波束成形网络 (基于注意力机制)
        self.beamforming = BeamformModel(**semconfig)
        
        # 信号传输
        self.transmitor = TransMitModel(self.Nc, self.Q)
        
        # ========== 信号解调模块 ========== #
        self.demodulator = DemodulateModel(self.Nc, self.Q, self.patches)
        
        # ========== 信道解码模块 ========== #
        self.Hdecode = nn.Sequential(
            nn.Linear(self.Nc, 512),
            nn.GELU(),
            nn.Linear(512, bsqconfig['embed_dim'])
        )
        
        self.image_quantizer.freeze_flows(encode=True, quantize=False, decode=False)
        
        self.chan_aware_loss = ChannelAwareLoss()
        self.sem_align_loss = SemanticAlignmentLoss()
    
    def get_last_layer(self):
        return self.image_quantizer.get_last_layer()

    def forward(self, pixel, H, H0, skip_quantize=False):
        """
        输入:
            pixel: [B*K, C, H, W]
            H: 信道矩阵 [B, K, Nc, Nt] (复数)
        输出:
            x_hat: 重建图像 [B*K, C, H, W]
        """
        B = pixel.shape[0]
        ### === 发送端 === ###
        # 1. BSQ图像量化编码
        quant, commit_loss, info = self.image_quantizer.encode(pixel, skip_quantize=skip_quantize) # [B*K, L, D=256]
        
        # 2. 联合信源信道编码
        symbol = self.Hadaptor(quant) # [B*K, L, Nc]
        
        # 3. qpsk 复数符号生成
        symbol = symbol.reshape(-1, self.Nc, self.patches)  # [B*K, Nc, L]
        transmitted = self.qpsk(symbol)  # [B*K, Nc, 2*Q]
        
        # 4. 信道状态信息提取
        H = H.reshape(B//self.K, self.K, self.Nc, 1, self.Nt)   # 每个batch下有K个用户，这K个用户由Nc*Nt个信道响应传输数据
        H0 = H0.reshape(B//self.K, self.K, self.Nc, 1, self.Nt)
        HH, noise, H_equ = self.csi_extractor(H, H0)
        
        # 5. 多用户波束成形（预编码） Beamforming
        signal = self.beamforming(transmitted, HH)
        
        # 6. 噪声注入
        power = (torch.sum(torch.abs(signal)**2,[2,3])).reshape(-1, self.Nc, 1, 1)
        noise = noise * torch.sqrt(power) / math.sqrt(self.Q)
        
        # 7. 信道传输
        received = self.transmitor.transmit(signal, noise, H_equ) # [B*K, Nc, 2*Q]
        chan_loss = self.chan_aware_loss(transmitted, received, H)
        
        ### === 接收端 === ###
        # 8. 信号解调
        symbol_hat = self.demodulator(received) #  [B*K, Nc, L]
        
        # 9. 信道解码
        symbol_hat = symbol_hat.reshape(-1, self.patches, self.Nc)  # [B*K, L, Nc]
        quant_hat = self.Hdecode(symbol_hat)  # [B*K, L, D]
        sem_loss = self.sem_align_loss(quant_hat, quant)
        
        # 10. 图像重建
        x_hat = self.image_quantizer.decode(quant)  # [B*K, C, H, W]
        
        return x_hat, (commit_loss, chan_loss, sem_loss), info
    
if __name__ == "__main__":
    from omegaconf import OmegaConf
    from transcoder.utils import config as config_utils
    from transcoder.data.comm_channel import SV_Channel, LAMP_GMMV

    cfg_file = f"{project_dir}/configs/connector/imagenet_bsqvit_l_snr_q.yaml"
    config = OmegaConf.load(cfg_file)
    model = config_utils.instantiate_from_config(config.model)
    model.cuda()

    B = config.data.batch_size
    semconfig = config.model.params.semconfig
    bsqconfig = config.model.params.bsqconfig
    K = semconfig.K
    C = 3
    H = bsqconfig.vitconfig.image_size
    W = bsqconfig.vitconfig.image_size
    Nc = semconfig.Nc
    Nt = semconfig.Nt
    paths = semconfig.paths
    pho = semconfig.pho
    L = semconfig.L
    SNR_dB = semconfig.SNR_dB
    
    image = torch.randn(B, C, H, W).cuda()
    H = SV_Channel(B, Nc, paths, [8,8], 1)
    # H0 = LAMP_GMMV(Nc, Nt, SNR_dB, K, L, pho, paths, H)
    H0 = H
    
    recon, loss, info = model(image, H, H0)
    
    print(recon.shape)
    print(loss)